import numpy as np
from opencv.img_utils import get_rgb_img, show_two_imgs

file_name = 'img/girl.png'


def median_filter(src, k_size):
    '中值滤波：选择矩阵内的中值填充当前像素  通常用于处理随机椒盐滤波效果较好'
    dst = np.zeros(src.shape, dtype=np.uint8)
    h = src.shape[0]
    w = src.shape[1]
    rx = k_size[0] // 2
    ry = k_size[1] // 2
    nc = src.shape[-1]
    for y in range(h):
        for x in range(w):
            n = 0
            v = np.empty(((2 * rx + 1) * (2 * ry + 1), nc), dtype=np.uint8)
            for i in range(-ry, ry + 1):
                iy = i + y
                if iy < 0 or iy >= h:
                    continue
                for j in range(-rx, rx + 1):
                    jx = j + x
                    if jx < 0 or jx >= w:
                        continue
                    v[n] = src[iy][jx]
                    n += 1
            v = np.resize(v, (n, nc))
            v.sort(axis=0)
            dst[y][x] = v[n // 2]

    return dst


def test():
    src = get_rgb_img(file_name)
    dst = median_filter(get_rgb_img(file_name), (4, 4))
    show_two_imgs(src, 'src', dst, 'dst')


test()
